package io.gatling.extensions.mysql.protocol

import java.util

import io.gatling.core.CoreComponents
import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.vertx.core.{Vertx, VertxOptions}
import io.vertx.mysqlclient.{MySQLConnectOptions, MySQLPool}
import io.vertx.sqlclient.PoolOptions

case class MysqlProtocol(host: String, port: Int = 3306, db: String, userName: String, passward: String) extends Protocol {}

object MysqlProtocol {
  val mysqlProtocolKey: ProtocolKey[MysqlProtocol, MysqlComponents] = new ProtocolKey[MysqlProtocol, MysqlComponents] {
    override def protocolClass: Class[Protocol] = classOf[MysqlProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): MysqlProtocol = MysqlProtocol(
      "127.0.0.1",
      3306,
      "test",
      "test_user",
      "test_password"
    )

    override def newComponents(coreComponents: CoreComponents): MysqlProtocol => MysqlComponents = mysqlProtocol => MysqlComponents(mysqlProtocol)
  }
}

case class MysqlComponents(mysqlProtocol: MysqlProtocol) extends ProtocolComponents {

  private val connectOptions: MySQLConnectOptions = new MySQLConnectOptions()
    .setHost(mysqlProtocol.host)
    .setPort(mysqlProtocol.port)
    .setDatabase(mysqlProtocol.db)
    .setUser(mysqlProtocol.userName)
    .setPassword(mysqlProtocol.passward)

  private val poolOptions: PoolOptions = new PoolOptions()
    .setMaxSize(1)

  private val verts: Vertx = Vertx.vertx(new VertxOptions())

  val clients: util.HashMap[Long, MySQLPool] = new util.HashMap[Long, MySQLPool]()

  override def onStart: Session => Session = session => {
    clients.put(session.userId, MySQLPool.pool(verts, connectOptions, poolOptions))
    session
  }

  override def onExit: Session => Unit = session => clients.get(session.userId).close()
}

case class MysqlProtocolBuilder(host: String, port: Int = 3306, db: String, userName: String, passward: String) {

  def build: MysqlProtocol = MysqlProtocol(host, port, db, userName, passward)
}

object MysqlProtocolBuilder {
  implicit def toMysqlProtocol(builder: MysqlProtocolBuilder): MysqlProtocol = builder.build
}
